import numpy as np
from numpy.matlib import repmat
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 步骤一（替换sans-serif字体）
plt.rcParams['axes.unicode_minus'] = False  # 步骤二（解决坐标轴负数的负号显示问题）
# 粒子群算法类
class PSO:
    def __init__(self):
        self.c1 = 2  # 学习因子
        self.c2 = 2  # 学习因子
        self.m = 100     # 粒子个数
        self.v = None
        self.x = None
        self.V_max = 0.01
        self.V_min = -0.01
        self.iteration_num = 100    # 迭代次数
        self.Pbest = None  # 局部最优点
        self.Gbest = None  # 全局最优点
        self.Gbest_Value = None  # 全局最优值
        self.Pbest_Value = None  # 局部最优值
        self.init_xy()  # 初始化粒子群

    # 目标函数
    def Camel(self, x, y):
        camel = (4 - 2.1 * x ** 2 + x ** 4 / 3) * x ** 2 + x * y + (-4 + 4 * y ** 2) * y ** 2
        return camel

    # 生成随机数组 10x2
    def init_xy(self):
        x = np.random.random((self.m, 1)) * 6 - 3
        y = np.random.random((self.m, 1)) * 4 - 2
        self.x = np.append(x, y, axis=1)
        self.v = np.random.random((self.m, 2)) * 0.001

    # 更新速度v、位置x
    def update_xv(self):
        self.v = self.v + self.c1 * np.random.random() * (self.Pbest - self.x) * 0.01 + self.c2 * np.random.random() * (
                self.Gbest - self.x) * 0.01
        self.x = self.x + self.v

    # 更新个体最优值
    def update_Pbest(self, Pose):
        if self.Pbest_Value is None:
            self.best_Value = Pose  # 最优值
            self.Pbest = self.x     # 最优值坐标点
        else:
            for i in range(self.m):
                # 若本次值小于局部历史最优值，则更新最优值和最优点
                if Pose[i] <= self.best_Value[i]:
                    # 更新最优值
                    self.best_Value[i] = Pose[i]
                    # 更新最优点
                    for j in range(2):
                        self.Pbest[i, j] = self.x[i, j]

    # 更新历史最优值
    def update_Gbest(self, Gbest, fx):
        # 第一次随机赋值
        if self.Gbest is None:
            self.Gbest = repmat(Gbest, self.m, 1)   # 坐标点
            self.Gbest_Value = fx                   # 函数值
        else:
            # 如果本次函数值小于全局最优值，则更新最优值，和最优点
            if fx <= self.Gbest_Value:
                # 更新最优值
                self.Gbest_Value = fx
                # 更新最优点
                self.Gbest[:, 0] = Gbest[0]
                self.Gbest[:, 1] = Gbest[1]
        return self.Gbest

    # 限制速度和位置
    def limit_xy(self):
        for i in range(self.m):
            # 限制Pose_x
            if self.x[i, 0] >=  3: self.x[i, 0] = 3
            if self.x[i, 0] <= -3: self.x[i, 0] = -3
            # 限制Pose_y
            if self.x[i, 1] >=  2: self.x[i, 1] = 2
            if self.x[i, 1] <= -2: self.x[i, 1] = -2

            # 限制V_x
            if self.v[i, 0] >= self.V_max: self.v[i, 0] = self.V_max
            if self.v[i, 0] <= self.V_min: self.v[i, 0] = self.V_min
            # 限制V_y
            if self.v[i, 1] >= self.V_max: self.v[i, 1] = self.V_max
            if self.v[i, 1] <= self.V_min: self.v[i, 1] = self.V_min

    # 解析
    def calculate(self):
        # 计算函数值
        pose = self.Camel(self.x[:, 0], self.x[:, 1])
        # 获得全局最优函数解
        fx_min = np.min(pose)
        # 获得索引值
        min_index = np.argmin(pose)
        # 获得最优值的（x , y)坐标值
        Gbest = self.x[min_index, :]
        # 更新全局最优解
        self.update_Gbest(Gbest, fx_min)
        # 更新个体最优解
        self.update_Pbest(pose)
        # 更新每个粒子的速度和位置
        self.update_xv()
        # 限制x,v
        self.limit_xy()
        return np.min(self.Gbest_Value)


    # 迭代函数
    def iteration(self):
        a = []
        for i in range(self.iteration_num):
            best = self.calculate()
            a.append(best)
        print('最优值：', self.Gbest_Value)
        print('最优点：', self.Gbest[0, :])
        return a

a = PSO()
f = a.iteration()
fig = plt.figure()
plt.xlabel('迭代次数')
plt.ylabel('最优函数值')
plt.plot(np.linspace(0, a.iteration_num, a.iteration_num), f, color="black", label='函数值')
plt.legend()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
xm = np.linspace(-3, 3, 100)
ym = np.linspace(-2, 2, 100)
xm, ym = np.meshgrid(xm, ym)
zm = a.Camel(xm, xm)
ax.scatter(a.Pbest[:, 0], a.Pbest[:, 1], np.min(a.Gbest_Value), color='r', s=20)
ax.plot_wireframe(xm, ym, zm, rstride=5, cstride=5, cmap=plt.cm.coolwarm)

plt.show()